from typing import Iterable, Union
import cv2
import matplotlib.pyplot as plt
import numpy as np

def rings(image_path : str, output_path : str, 
        blur_size : int = 1, canny_up_threshold : int = 170, bias : int|str = 'auto',
        plot_dict : dict = {"color" : "k", "linewidth" : 0.5}, 
        default_height : float = 50/6, layers_gate = [5000,3000,1200,400]) -> None:

    layers_gate.insert(0,1e10)
    # read image
    image = cv2.imread(image_path)
    image = cv2.rotate(image,cv2.ROTATE_180)
    width = image.shape[1]
    height = image.shape[0]
    if width > 2400 :
        image = cv2.resize(image,(2400,2400*height//width))
        width = image.shape[1]
        height = image.shape[0]
    true_ratio = width / height
    # regular operation, no more explainations
    
    blur = cv2.GaussianBlur(image, (blur_size,blur_size), 0)
    edge = cv2.Canny(blur, 0, canny_up_threshold)
    # sharpen_op = np.array([[0, -1, 0], [-1, 5, -1], [0, -1, 0]], dtype=np.float32)
    # sharpen_image = cv2.filter2D(blur, cv2.CV_8UC1, sharpen_op)
    # sharpen_image = cv2.filter2D(sharpen_image, cv2.CV_8UC1, sharpen_op)
    # gray = cv2.cvtColor(sharpen_image, cv2.COLOR_BGR2GRAY)
    # ret, thresh = cv2.threshold(gray, 127, 255, cv2.THRESH_BINARY)

    # get contour of the edge image
    contour_tuple = cv2.findContours(edge, mode=cv2.RETR_TREE, method=cv2.CHAIN_APPROX_NONE)
    contours = list(contour_tuple[0])

    if bias == 'auto':
        bias = (height*width)**0.5/700

    for i in range(len(contours)):
        ring = contours[i]
        tmp = cv2.approxPolyDP(ring,bias,False)
        contours[i] = tmp
    
    contours_filtered = []
    length_list = []

    for i in range(len(contours)):
            ring = contours[i]
            length = 0
            x = ring[0,0,0]
            y = ring[0,0,1]
            for k in ring:
                length += ((x-k[0,0])**2 + (y-k[0,1])**2)**0.5 / width * 600
            length_list.append(length)
    for p in range(len(layers_gate)-1):
        for i in range(len(contours)):
            ring = contours[i]
            length = length_list[i]
            if length > layers_gate[p+1] and length <= layers_gate[p]:
                contours_filtered.append(ring)
        pass
    contours = contours_filtered

    rings = [np.array(c).reshape([-1, 2]) for c in contours]

    # adjust coordinate system to the image coordinate system
    max_x, max_y, min_x, min_y = 0, 0, width, height

    # adjust ratio
    plt.figure(figsize=[default_height * true_ratio, default_height])

    # plot to the matplotlib
    for _, ring in enumerate(rings):
        close_ring = np.vstack((ring, ring[0]))
        xx = close_ring[..., 0]
        yy = max_y - close_ring[..., 1]
        plt.plot(xx, yy, **plot_dict, )

    plt.axis("off")
    plt.tight_layout()
    plt.xlim(0,width)
    plt.ylim(-height,0)
    plt.savefig(output_path)
    plt.savefig(output_path+'.png')
    #plt.show()
    pass

if __name__ == '__main__':
    rings('out.jpg', './img_canny.svg')